import re
from c2rust.actions import *
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.logs import logger
from rag.indexing_script import get_service_context, init_query_engine, QUERY_ENGINE

def get_class_name(cls) -> str:
    """Return class name"""
    return f"{cls.__module__}.{cls.__name__}"

class RustAnalyzer(Role):
    name: str = "RustAnalyzer"
    profile: str = "RustAnalyzer"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([AddMainFunc, CorrectCode])
        self.set_actions([AnalyzeCode, AnalyzeCodeAfterCorrect])

    async def _think(self):
        news = self.rc.news[0]
        logger.info(news.cause_by)
        logger.info(type(news.cause_by))
        logger.info(get_class_name(CorrectCode))
        if news.cause_by == get_class_name(CorrectCode):
            self.rc.todo = AnalyzeCodeAfterCorrect()
        else:
            self.rc.todo = AnalyzeCode()
        logger.info(self.rc.todo)
        
    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg

class RustTutorialRetriever(Role):
    name: str = "RustTutorialRetriever"
    profile: str = "RustTutorialRetriever"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([AnalyzeCode, AnalyzeCodeAfterCorrect])
        self.set_actions([RetrieveTutorial, RetrieveTutorialAfterCorrect])
        

    async def _think(self):
        news = self.rc.news[0]
        logger.info(news.cause_by)
        logger.info(type(news.cause_by))
        logger.info(get_class_name(AnalyzeCodeAfterCorrect))
        if news.cause_by == get_class_name(AnalyzeCodeAfterCorrect):
            self.rc.todo = RetrieveTutorialAfterCorrect()
        else:
            self.rc.todo = RetrieveTutorial()
        logger.info(self.rc.todo)


    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(context[0].content, QUERY_ENGINE)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg


class RustEnhancerRAG(Role):
    name: str = "RustEnhancerRAG"
    profile: str = "RustEnhancerRAG"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([RetrieveTutorial, RetrieveTutorialAfterCorrect])
        self.set_actions([EnhanceCodeRAG, EnhanceCodeAfterCorrectRAG])
        
    async def _think(self):
        news = self.rc.news[0]
        logger.info(news.cause_by)
        logger.info(type(news.cause_by))
        logger.info(get_class_name(RetrieveTutorialAfterCorrect))
        if news.cause_by == get_class_name(RetrieveTutorialAfterCorrect):
            self.rc.todo = EnhanceCodeAfterCorrectRAG()
        else:
            self.rc.todo = EnhanceCodeRAG()
        logger.info(self.rc.todo)

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg
